# -*- coding: utf-8 -*-

#    OriginDB
#    Copyright (C) 2010  Jason Power <power.jg@gmail.com>
#
#    Algorithm Copyright 2003-2008 Beda Kosata <beda@zirael.org>
#    OASA project (http://bkchem.zirael.org/oasa_en.html)
#
#    This program is free software; you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation; either version 2 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License along
#    with this program; if not, write to the Free Software Foundation, Inc.,
#    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

import openbabel

from math import pi, sqrt, cos, sin, atan2




class GenCoords(object):
	def __init__(self, bondLength=1.0):
		self.bondLength = bondLength

		self.coords = {}
		self.mol = None

	def setCoords(self):
		if not self.mol or not self.coords:
			return

		for atom in openbabel.OBMolAtomIter(self.mol):
			c = self.coords[atom.GetIdx()]
			atom.SetVector(c[0], c[1], 0)

	def calcCoords(self, mol, force=0):

		processed = []
		self.mol = mol

		# Dict keyed on atom idx's with a list of the position as the value
		self.coords = {}

		# we create a backbone, it is either one of the rings or
		# the longest chain in case of acyclic molecule
		self.rings = mol.GetSSSR()
		self.rings = [r for r in self.rings]

		if self.rings:
			# find the most crowded ring
			jmax = 0
			imax = 0
			for i, ring in enumerate(self.rings):
				j = 0
				for r in self.rings:
					if r != ring and (set(r._path) & set(ring._path)):
						j += 1
				if j > jmax:
					jmax = j
					imax = i

			backbone = self.rings[imax]
			del self.rings[imax]

			gcoords = genRingCoords(len(backbone._path))
			for idx in backbone._path:
				self.coords[idx] = gcoords.next()
				processed.append(idx)

			processed += self.processAllAnelatedRings(backbone._path)

		elif mol.NumAtoms == 1:
			# if there is only 1 atom then the Idx is 1
			self.coords[1] = [0,0]
			processed.append(1)

		else:
			# chain - we process 2 atoms as backbone and leave the rest for the code
			a1 = openbabel.OBMolAtomIter(self.mol).next()
			bond = openbabel.OBAtomBondIter(a1).next()
			a2 = bond.GetEndAtom()
			if a1.GetIdx() == a2.GetIdx():
				a2 = bond.GetBeginAtom()

			self.coords[a1.GetIdx()] = [0,0]
			processed.append(a1.GetIdx())
			self.coords[a2.GetIdx()] = [self.bondLength, 0]
			processed.append(a2.GetIdx())

			self.rings = []

		self.continueProcessing(processed)


	def continueProcessing(self, processed):
		"""processes the atoms in circles around the backbone (processed) until all is done"""
		while processed:
			newProcessed = []
			for v in processed:
				if len(self.coords) == self.mol.NumAtoms():
					# its all done
					return
				# look if v is part of a ring
				ring = None
				for r in self.rings:
					if v in r._path:
						ring = r
						break

				if not ring:
					# v is not in a ring - we can continue
					newProcessed += self.processAtomNeighbors( v)

				else:
					# v is in ring so we process the ring
					if len(processed) > 1 and set(processed) <= set(ring._path):
						# not sure if this is right at all...
						print "I don't know if this is right, hopefully it won't be called"
						newProcessed += self.processAllAnelatedRings(processed)
					else:
						del self.rings[self.rings.index(ring)]
						ring = self.sortRing(ring, startFrom=v)
						del ring[0]

						a = None
						for atom in openbabel.OBMolAtomIter(self.mol):
							if atom.GetIdx() == v:
								a = atom
								break

						d = None
						for bond in openbabel.OBAtomBondIter(a):
							end = bond.GetEndAtomIdx()
							start = bond.GetBeginAtomIdx()
							if end != v and end in self.coords:
								d = end
							elif start != v and start in self.coords:
								d = start

						ca = clockwiseAngleFromEast(self.coords[v][0]-self.coords[d][0], self.coords[v][1]-self.coords[d][1])
						size = len(ring) +1
						da = degToRad(180 -180*(size-2)/size)
						gcoords = genAngleStream(da, startFrom=ca -pi/2+da/2)

						# here we generate the coords
						self.applyGenToAtoms(gcoords, ring, v)
						ring.append(v)
						newProcessed += ring
						newProcessed += self.processAllAnelatedRings(ring)

			processed = newProcessed


	def processAllAnelatedRings(self, baseIdxs):
		out = []
		toGo = []

		for r in self.rings:
			if set(r._path) & set(baseIdxs):
				toGo.append(r)

		for r in toGo:
			out += self.processOneAnelatedRing(baseIdxs)
		if out:
			for b in toGo:
				out += self.processAllAnelatedRings(b._path)

		return out


	def processOneAnelatedRing(self, baseIdxs):
		out = []
		ring = None
		i = 0
		for r in self.rings:
			if set(r._path) & set(baseIdxs):
				ring = r
				break
			i += 1
		if ring:
			del self.rings[i]
			inter = set(ring._path) & set(baseIdxs)
			# add idx if not already done...
			toGo = [idx for idx in ring._path if idx not in self.coords]
			if (len(ring._path) - len(toGo)) == len(inter):
				out += self.processSimpleAnelatedRing(ring._path, baseIdxs)
			else:
				out += self.processMultiAnelatedRing(ring._path)
		return out


	def processAtomNeighbors(self, v):
		def getAngleAtSide(v, d, d2, relation, attachAngle):
			if attachAngle == 180:
				return attachAngle
			side = onWhichSideIsPoint( (self.coords[d][0], self.coords[d][1], self.coords[v][0], self.coords[v][1]),
									   (self.coords[d2][0], self.coords[d2][1]))
			an = angle + degToRad(attachAngle)
			x = self.coords[v][0] + self.bondLength*cos(an)
			y = self.coords[v][1] + self.bondLength*sin(an)
			if relation*side == onWhichSideIsPoint((self.coords[d][0], self.coords[d][1], self.coords[v][0], self.coords[v][1]),
												   (x,y)):
				return attachAngle
			else:
				return -attachAngle

		baseNeighs = getNeighbors(v, self.mol)

		toGo = [a for a in baseNeighs if a not in self.coords]
		done = [a for a in baseNeighs if a not in toGo]

		if len(done) == 1 and (len(toGo) == 1):
			# only simple non-branched chain
			d = done[0]
			t = toGo[0]

			#decide angle
			angleToAdd = 120
			# triple bonds
			bond = getBond(t, v, self.mol)
			if bond.GetBO() == 3:
				angleToAdd = 180
			# two double bonds in a row
			elif bond.GetBO() == 2:
				# get bo of previous thing
				bond2 = getBond(d, v, self.mol)
				if bond2.GetBO() == 2:
					angleToAdd = 180

			angle = clockwiseAngleFromEast(self.coords[d][0]-self.coords[v][0], self.coords[d][1]-self.coords[v][1])

			dns = getNeighbors(d, self.mol)
			if len(dns) == 2:
				d2 = (dns[0] == v) and dns[1] or dns[0]
				if d2 in self.coords:
					angleToAdd = getAngleAtSide(v, d, d2, -1, angleToAdd)

			an = angle + degToRad(angleToAdd)
			self.coords[t] = [self.coords[v][0] + self.bondLength*cos(an), self.coords[v][1] + self.bondLength*sin(an)]

		elif toGo:
			# branched chain
			angles = [clockwiseAngleFromEast(self.coords[at][0] - self.coords[v][0], self.coords[at][1] - self.coords[v][1]) for at in done]
			angles.append(2*pi + min(angles))
			angles.sort(reverse=True)

			diffs = listDifferences(angles)
			i = diffs.index( max(diffs))
			angle = (angles[i] - angles[i+1])/(len(toGo)+1)
			gcoords = genCoordsFromStream(genAngleStream(angle, startFrom=angles[i+1]+angle), length = self.bondLength)

			for a in toGo:
				dx, dy = gcoords.next()
				self.coords[a] = [self.coords[v][0]+dx, self.coords[v][1]+dy]

		return toGo


	def applyGenToAtoms( self, gen, atoms, start, bondLength=None):
		bl = bondLength or self.bondLength
		x, y = self.coords[start][0], self.coords[start][1]
		for v in atoms:
			a = gen.next()
			x += bl*cos( a)
			y += bl*sin( a)
			self.coords[v] = [x, y]

	def processSimpleAnelatedRing(self, ring, base):
		out = []
		inter = [idx for idx in ring if idx in self.coords]
		if len(inter) == 1:
			# rings are connected via one atom
			idx = inter[0]
			ring = self.sortRing(ring, startFrom=idx)

			baseNeighs = getNeighbors(idx, self.mol)

			if len(baseNeighs) < 2:
				raise "should have more neighbors!!"

			d1 = baseNeighs[0]
			d2 = baseNeighs[1]

			ca1 = clockwiseAngleFromEast(self.coords[idx][0] - d1.x, self.coords[idx][1] - d1.y)
			ca2 = clockwiseAngleFromEast(self.coords[idx][0] - d2.x, self.coords[idx][1] - d2.y)

			ca = (ac1 + ca2) / 2.0
			if abs(ca1-ca2) < pi:
				ca += -pi/2
			else:
				ca += pi/2

			size = len(ring)
			da = degToRad(180 - 180*(size-2)/size)
			gcoords = genAngleStream(da, startFrom=ca+da/2)
			del ring[0]
			self.applyGenToAtoms(gcoords, ring, idx)
			ring.append(idx)
			out += ring
		elif len(inter) == 2:
			# There are two common atoms
			v1, v2 = inter
			ring = self.sortRing(ring, startFrom=v1)
			del ring[0]
			if ring[0] == v2:
				del ring[0]
				# v1 should be neighboring the new ring[0]
				v1, v2 = v2, v1
			else:
				del ring[-1]

			side = sum([onWhichSideIsPoint((self.coords[v1][0],self.coords[v1][1],self.coords[v2][0],self.coords[v2][1]),
										  (self.coords[v][0],self.coords[v][1])) for v in base])
			if not side:
				raise "this shouldn't happen!"

			ca = clockwiseAngleFromEast(self.coords[v1][0]-self.coords[v2][0], self.coords[v1][1]-self.coords[v2][1])
			size = len(ring)+2
			da = degToRad(180 - 180*(size-2)/size)
			if side > 0:
				da = -da

			gcoords = genAngleStream(da, startFrom=ca+da)
			self.applyGenToAtoms(gcoords, ring, v1)
			ring.append(v1)
			ring.append(v2)
			out += ring
		else:
			# there are more than 2 atoms common
			if len(ring) == len(base):
				out += self.processMultiAnelatedRing(ring, angleShift=15)
			else:
				out += self.processMultiAnelatedRing(ring)

		return out

	def processMultiAnelatedRing(self, ring, angleShift=0):
		out = []
		toGo = [v for v in ring if v not in self.coords]
		if not toGo:
			return out

		back = [v for v in ring if v in self.coords]

		# I'm not sure that this is right.
		sortedBack = []
		for i in ring:
			if i in back:
				sortedBack.append(i)

		# sort togo??

		v1 = sortedBack[0]
		v2 = sortedBack[-1]
		v3 = sortedBack[1]

		i = ring.index(toGo[0])
		if i == len(ring) - 1:
			if ring[0] != v1 and ring[i-1] != v1:
				v1, v2 = v2, v1
		else:
			if ring[i-1] != v1 and ring[i+1] != v1:
				v1, v2 = v2, v1

		blockedAngle = sumOfRingInternalAngles(len(back))
		overallAngle = sumOfRingInternalAngles(len(ring))

		da = optimalRingInternalAngle(len(ring))

		#if there are 2 rings of same size inside each other, we need to use the angle_shift
		if angleShift:
			da += 2*angleShift/len(toGo)

		ca = degToRad( 180-(overallAngle - blockedAngle - len( toGo) * da + angleShift)/2)  # connection angle

		side = sum([onWhichSideIsPoint((self.coords[v1][0],self.coords[v1][1],self.coords[v2][0],self.coords[v2][1]),
									  (self.coords[v][0],self.coords[v][1])) for v in back if v != v1 and v != v2])

		# we need to make sure that the ring is drawn on the right side
		if side > 0:
			ca = -ca

		ca += clockwiseAngleFromEast(self.coords[v1][0]-self.coords[v2][0], self.coords[v1][1]-self.coords[v2][1])

		da = 180 - da # for drawing we use external angle

		# we must ensure that the ring will progress towards the second end
		if onWhichSideIsPoint((self.coords[v1][0],self.coords[v1][1],self.coords[v3][0],self.coords[v3][1]),
							  (self.coords[v2][0],self.coords[v2][1])) < 0:
			da = -da

		# dry run to see where we get
		gcoords = genAngleStream(degToRad(da), startFrom=ca)
		x, y = self.coords[v1][0], self.coords[v1][1]
		for i in range(len(toGo)+1):
			a = gcoords.next()
			x += self.bondLength*cos( a)
			y += self.bondLength*sin( a)
		# end of dry run, we can scale the bond_length now
		length = sqrt((self.coords[v1][0]-self.coords[v2][0])**2 + (self.coords[v1][1]-self.coords[v2][1])**2)
		realLength = sqrt( (self.coords[v1][0] - x)**2+ (self.coords[v1][1] - y)**2)
		bl = self.bondLength * length / realLength
		gcoords = genAngleStream( degToRad( da), startFrom= ca)
		# and here we go
		self.applyGenToAtoms( gcoords, toGo, v1, bondLength=bl)
		out += toGo
		return out


	def sortRing(self, ring, startFrom=None):
		r = []
		try:
			ring = ring._path
		except AttributeError:
			pass

		for idx in ring:
			r.append(idx)
		ans = []
		ans.append(startFrom)

		si = r.index(startFrom)
		i = si + 1
		while i != si:
			if i >= len(r):
				i -= len(r)
				continue
			ans.append(r[i])
			i += 1

		return ans

def sumOfRingInternalAngles(size):
	return (size-2)*180

def optimalRingInternalAngle(size, angleShift=0):
	return sumOfRingInternalAngles(size)/size

def genRingCoords(size, length=1):
	coords = genCoordsFromDegStream(genAngleStream(180-180*(size-2)/size),length)

	x, y = 0, 0
	for i in range(size):
		dx, dy = coords.next()
		x += dx
		y += dy
		yield x, y

def genAngleStream(angle, startFrom=0, alternate=0):
	yield startFrom

	if alternate:
		while 1:
			yield angle
			yield -angle
	else:
		a = startFrom
		while 1:
			a += angle
			yield a

def genCoordsFromDegStream(stream, length=1):
	for a in stream:
		ang = degToRad(a)
		yield (length*cos(ang), length*sin(ang))

def genCoordsFromStream(stream, length=1):
	for a in stream:
		yield (length*cos(a), length*sin(a))

def degToRad(deg):
	return pi*deg/180

def radToDeg(rad):
	return 180*rad/pi

def clockwiseAngleFromEast(x, y):
	angle = atan2(y, x)
	if angle < 0:
		angle = 2*pi + angle
	return angle

def onWhichSideIsPoint(line, point):
	""" tells whether a point is on one side of a
		line or on the other (1,0,-1) - 0 is for point on line.
		line is given as sequence of four coordinates, point
		as sequence of two coords
	"""

	x1, y1, x2, y2 = line
	x, y = point

	a = atan2(y-y1, x-x1)
	b = atan2(y2-y1, x2-x1)

	if a*b < 0 and abs(a-b) > pi:
		if a < 0:
			a += 2*pi
		else:
			b += 2*pi

	if abs(a-b) <= 0 or abs(abs(a-b) -pi) <= 0:
		return 0
	if a-b < 0:
		return 1
	else:
		return -1


def listDifferences(l):
	ret = []
	for i in range( len( l)-1):
		ret.append( l[i]-l[i+1])
	return ret

def inRing(idx, mol):
	atom = None
	for a in openbabel.OBMolAtomIter(mol):
		if a.GetIdx() == idx:
			atom = a
			break
	if not atom:
		raise "huh this should be impossible"

	return atom.MemberOfRingCount()

def getAtom(idx, mol):
	atom = None
	for a in openbabel.OBMolAtomIter(mol):
		if a.GetIdx() == idx:
			atom = a
			break
	if not atom:
		raise "huh this should be impossible"

	return atom

def getBond(idx1, idx2, mol):
	a1 = getAtom(idx1, mol)
	a2 = getAtom(idx2, mol)

	b = mol.GetBond(a1, a2)
	if not b:
		b = mol.GetBond(a1, a2)

	if not b:
		raise "There should be a bond here..."

	return b

def getNeighbors(idx, mol):
	atom = getAtom(idx, mol)

	baseNeighs = []
	for bond in openbabel.OBAtomBondIter(atom):
		if bond.GetEndAtomIdx() != idx:
			baseNeighs.append(bond.GetEndAtomIdx())
		else:
			baseNeighs.append(bond.GetBeginAtomIdx())

	return baseNeighs